dirtools.py 1.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. import os
  2. import io
  3. import contextlib
  4. import tempfile
  5. import shutil
  6. import errno
  7. import zipfile
  8. @contextlib.contextmanager
  9. def tempdir():
  10. """Create a temporary directory in a context manager."""
  11. td = tempfile.mkdtemp()
  12. try:
  13. yield td
  14. finally:
  15. shutil.rmtree(td)
  16. def mkdir_p(*args, **kwargs):
  17. """Like `mkdir`, but does not raise an exception if the
  18. directory already exists.
  19. """
  20. try:
  21. return os.mkdir(*args, **kwargs)
  22. except OSError as exc:
  23. if exc.errno != errno.EEXIST:
  24. raise
  25. def dir_to_zipfile(root):
  26. """Construct an in-memory zip file for a directory."""
  27. buffer = io.BytesIO()
  28. zip_file = zipfile.ZipFile(buffer, 'w')
  29. for root, dirs, files in os.walk(root):
  30. for path in dirs:
  31. fs_path = os.path.join(root, path)
  32. rel_path = os.path.relpath(fs_path, root)
  33. zip_file.writestr(rel_path + '/', '')
  34. for path in files:
  35. fs_path = os.path.join(root, path)
  36. rel_path = os.path.relpath(fs_path, root)
  37. zip_file.write(fs_path, rel_path)
  38. return zip_file